#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Copyright © 启明星辰 版权所有
# @Time    : 2024/5/13 14:27
# @Author  : 胡浩浩
# @File    : mitmweb_start_old.py
# @IDE     : PyCharm
# @description :
import base64
import json
import os
import random
import time

import jmespath
import uvicorn
from pywebio import config, session
from pywebio.output import *
from pywebio.pin import put_input, pin
from pywebio.platform.tornado_http import start_server
from pywebio.session import run_js, defer_call, eval_js, download
from typing import Any
import docker
import ruamel.yaml
from pywebio.platform.fastapi import asgi_app

from pywebio.session import defer_call
from fastapi import FastAPI
from starlette.responses import RedirectResponse

app = FastAPI()


@app.get("/")
async def redirect():
    # 重定向到指定的URL
    return RedirectResponse(url="/search")


css = """
.container {
    margin-top: 0;
    max-width: 60%;
}
"""
config(title="测试用例转化", theme="minty", css_style=css)

listen_port = 8181
web_port = 8081

# 创建 Docker 客户端对象
client = docker.from_env()


def read_yml(file_path, node=None):
    with open(file_path, "r", encoding="utf-8") as f:
        yaml = ruamel.yaml.YAML(typ='unsafe', pure=True)
        yaml.allow_duplicate_keys = True
        data = yaml.load(f)
    if node:
        return data[node]
    else:
        return data if data else []


def write_to_yml(file_path, data):
    # 顺序写入新的yml文件中,保证格式不混乱
    with open(file_path, "w", encoding="utf-8") as f:
        ruamel.yaml.round_trip_dump(data, f, default_flow_style=False, allow_unicode=True)


def stop_and_restart_docker_container(container_id, filter_express=None):
    try:
        # 停止容器
        container = client.containers.get(container_id)
        container.start()
        print(f"Container {container_id} stopped successfully.")
        # 复制文件到容器内指定目录
        with open("../src/mitmproxy_plugin_temp.py", 'rb') as f:
            data = f.read().decode('utf-8')
            data = data.replace("{{filter_express}}", filter_express)
            # print(data)
        with open("../src/source/mitmproxy_plugin.py", 'wb') as f:
            f.write(data.encode('utf-8'))
        # 重启容器
        # 构建 docker cp 命令
        docker_cp_command = f"docker cp mitmproxy_plugin.py {container_id}:/home/mitmproxy"

        # 使用 os.system 执行命令
        os.system(docker_cp_command)
        print(f"File 'mitmproxy_plugin_temp.py' copied to container '{container_id}' successfully.")
        container.start()
        print("File copied successfully and container restarted.")
        return True
    except docker.errors.NotFound as e:
        print(f"Container {container_id} not found: {e}")
        return False
    except docker.errors.APIError as e:
        print(f"Error stopping or removing container {container_id}: {e}")
        return False


def start_mitmproxy_container(listen_port: int = 8181, web_port: int = 8081) -> tuple[Any, Any]:
    # 容器配置
    container_config = {
        'image': 'ptesting_mitmproxy',  # 镜像名称
        'command': [
            'mitmweb',
            '-s', 'mitmproxy_plugin.py',
            '--set', 'ssl_insecure=true',
            '--listen-host', '0.0.0.0',
            '--listen-port', f'{listen_port}',
            '--web-host', '0.0.0.0',
            '--web-port', f'{web_port}'
        ],
        "network_mode": "host",
        'detach': True,
        'tty': True
    }

    # 启动容器
    container = client.containers.run(**container_config)

    container = client.containers.get(container.id)
    # 返回容器 ID
    # 5d688e6cba9e
    return container.id[:12]


def show_result(title, container_id, search_kw, data_containers_list):
    listen_port = jmespath.search(f'@[?@[0]==`{container_id}`][1]|[0]', data_containers_list)
    web_port = jmespath.search(f'@[?@[0]==`{container_id}`][2]|[0]', data_containers_list)
    stop_and_restart_docker_container(container_id, search_kw)
    return popup('运行结果', [
        put_html(f'<h3>{title}</h3>'),
        put_html('<p>代理IP：192.168.114.37</p>'),
        put_html(f'<p>代理端口：{listen_port}</p>'),
        put_row([
            put_html(f'<p>web地址：http://192.168.114.37:{web_port}</p>'),
            put_button("跳转", onclick=lambda: eval_js(
                f'window.open("http://192.168.114.37:{web_port}")'), small=True)
        ], size="60%"),
        put_row([
            None,
            put_button("关闭", onclick=lambda: close_popup(), color="dark")
        ], size="80%")
    ], closable=False)


def runs(user_ip=None):
    """
        业务逻辑：
        用data_containers_list.yml存储资源相关信息
        用data_containers_user_relations_list.yml存储用户与资源的关系
        1.判断搜索框有没有值，如果没有直接结束程序
        2.判断用户ip有没有分配过资源，如果已经分配过资源了，使用已经分配的资源，如果没有分配过资源，则需要新分配资源
        3.分配资源时需要验证资源的数量，如果超过10个则不允许新增资源，直接从原有的资源中那些没有被用到的资源，进行重新分配
        4.用户断开连接时需要释放资源
    """
    global listen_port
    global web_port
    # 验证搜索条件不为空
    search_kw = pin.search
    print("我拿到了查询条件", search_kw)
    if not search_kw:
        toast("查询条件不允许为空，请输入测试环境host或者api条件", color="red")
        return
    print("我拿到了查询条件", user_ip)
    data_containers_list = read_yml("../src/data_containers_list.yml")
    if len(data_containers_list) == 0:
        for i in range(5):
            new_container_id = start_mitmproxy_container(listen_port + i, web_port + i)
            data_containers_list.append([new_container_id, listen_port + i, web_port + i])
            write_to_yml("../src/data_containers_list.yml", data_containers_list)

    data_containers_user_relations_list = read_yml("../src/data_containers_user_relations_list.yml")
    print("我拿到的", len(data_containers_user_relations_list))
    print("我拿到的", data_containers_user_relations_list)
    container_id = jmespath.search(f'@[?@[0]==`{user_ip}`][1]|[0]', data_containers_user_relations_list)
    if len(data_containers_user_relations_list) >= 5 and not container_id:
        toast("最多支持5人同时在线，请排队等待", color="red")
        return
    else:
        if container_id:
            show_result("服务启已重启，请继续使用之前打开的web地址", container_id, search_kw, data_containers_list)
        else:
            # 直接从已有的资源池里面去重新分配资源
            all_containers_id_list = jmespath.search('@[*][0]', data_containers_list)
            used_containers_id_list = jmespath.search('@[*][1]', data_containers_user_relations_list)
            # 将列表转换为集合，并计算差集
            difference = list(set(all_containers_id_list) - set(used_containers_id_list))
            print("拿到了可以使用的资源")
            print(difference)
            container_id = random.choice(difference)
            print(container_id)
            if container_id:
                show_result("已为你分配代理资源", container_id, search_kw, data_containers_list)
                data_containers_user_relations_list.append([user_ip, container_id, int(time.time())])
                write_to_yml("../src/data_containers_user_relations_list.yml", data_containers_user_relations_list)


def reload():
    js_code = "history.go(0);"
    run_js(js_code)


def downloads(file_name):
    with open(f".mitmproxy/{file_name}", mode="rb") as f:
        file_content = f.read()
    # 下载文件
    download(file_name, file_content)


def show():
    user_ip = session.info.user_ip
    print("代理", session.info.user_agent)
    put_row([
        put_html("<h3>mitmproxy测试用例转化</h3>")
    ])
    put_html("<h4>第一次使用请先安装对应操作系统的证书</h4>")
    put_row([
        put_buttons(
            ["Windows", "Linux", "macOS", "Android", "iOS", "Firefox", "Other Platforms"],
            onclick=[lambda:downloads("mitmproxy-ca-cert.p12"),
                     lambda:downloads("mitmproxy-ca-cert.pem"),
                     lambda:downloads("mitmproxy-ca-cert.pem"),
                     lambda:downloads("mitmproxy-ca-cert.cer"),
                     lambda:downloads("mitmproxy-ca-cert.pem"),
                     lambda:downloads("mitmproxy-ca-cert.pem"),
                     lambda:downloads("mitmproxy-ca-cert.pem")], small=True),
    ])
    put_row([
        put_text("产品地址 *"),
        put_input("search", placeholder="请输入产品地址,如：https://192.168.114.91:8443/csa"),
        None,
        put_button("开始", onclick=lambda: runs(user_ip)),
    ], size="10% 60% 10%")

    @defer_call
    def cleanup():
        data_containers_user_relations_list = read_yml("../src/data_containers_user_relations_list.yml")
        current_user_relation = jmespath.search(f"@[?@[0]==`{user_ip}`][]", data_containers_user_relations_list)
        if current_user_relation:
            data_containers_user_relations_list.remove(current_user_relation)
            write_to_yml("../src/data_containers_user_relations_list.yml", data_containers_user_relations_list)
            container = client.containers.get(current_user_relation[1])
            container.stop()


app.mount("/search", asgi_app(show, cdn=False))

if __name__ == '__main__':
    uvicorn.run(app, port=7777, host="0.0.0.0")
